{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 04: computation of the inf-sup constant for a Stokes problem discretization\n",
    "\n",
    "In this tutorial we compare the computation of the inf-sup constant of a Stokes problem using standard assembly with mixed function spaces and block assembly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "from slepc4py import SLEPc\n",
    "from ufl import (div, dx, FiniteElement, grad, inner, MixedElement, split, TestFunction, TrialFunction,\n",
    "                 VectorElement)\n",
    "from dolfinx import Constant, Function, FunctionSpace, UnitSquareMesh\n",
    "from dolfinx.cpp.fem import create_vector_block\n",
    "from dolfinx.fem import locate_dofs_topological\n",
    "from dolfinx.mesh import locate_entities_boundary\n",
    "from dolfinx.plot import create_vtk_topology\n",
    "from multiphenicsx.fem import (assemble_matrix, assemble_matrix_block, assemble_scalar, BlockVecSubVectorWrapper,\n",
    "                               DofMapRestriction, VecSubVectorWrapper)\n",
    "import pyvista"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = UnitSquareMesh(MPI.COMM_WORLD, 32, 32)\n",
    "\n",
    "\n",
    "def wall(x):\n",
    "    return np.logical_or(x[1] < 0 + np.finfo(float).eps, x[1] > 1 - np.finfo(float).eps)\n",
    "\n",
    "\n",
    "boundary_facets = locate_entities_boundary(mesh, mesh.topology.dim - 1, wall)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dolfinx_to_pyvista_mesh(mesh):\n",
    "    num_cells = mesh.topology.index_map(mesh.topology.dim).size_local\n",
    "    cell_entities = np.arange(num_cells, dtype=np.int32)\n",
    "    pyvista_cells, cell_types = create_vtk_topology(mesh, mesh.topology.dim, cell_entities)\n",
    "    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)\n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_mesh_plot(mesh):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_mesh_plot(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V_element = VectorElement(\"Lagrange\", mesh.ufl_cell(), 2)\n",
    "Q_element = FiniteElement(\"Lagrange\", mesh.ufl_cell(), 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary function for eigenvector normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(u1, u2, p):\n",
    "    u1_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u1), grad(u1)) * dx), op=MPI.SUM))\n",
    "    u1.vector.scale(1. / u1_norm)\n",
    "    u1.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)\n",
    "    u2_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u2), grad(u2)) * dx), op=MPI.SUM))\n",
    "    u2.vector.scale(1. / u2_norm)\n",
    "    u2.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)\n",
    "    p_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(p * p * dx), op=MPI.SUM))\n",
    "    p.vector.scale(1. / p_norm)\n",
    "    p.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard FEniCSx formulation using a mixed function space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_monolithic():\n",
    "    # Function spaces\n",
    "    W_element = MixedElement(V_element, Q_element)\n",
    "    W = FunctionSpace(mesh, W_element)\n",
    "\n",
    "    # Test and trial functions: monolithic\n",
    "    vq = TestFunction(W)\n",
    "    (v, q) = split(vq)\n",
    "    up = TrialFunction(W)\n",
    "    (u, p) = split(up)\n",
    "\n",
    "    # Variational forms\n",
    "    lhs = inner(grad(u), grad(v)) * dx - div(v) * p * dx - div(u) * q * dx\n",
    "    rhs = - inner(p, q) * dx\n",
    "\n",
    "    # Define restriction for DOFs associated to homogenous Dirichlet boundary conditions\n",
    "    dofs_W = np.arange(0, W.dofmap.index_map.size_local + W.dofmap.index_map.num_ghosts)\n",
    "    bdofs_V = locate_dofs_topological(W.sub(0), mesh.topology.dim - 1, boundary_facets)\n",
    "    bdofs_V = locate_dofs_topological((W.sub(0), W.sub(0).collapse()), mesh.topology.dim - 1, boundary_facets)[0]\n",
    "    restriction = DofMapRestriction(W.dofmap, np.setdiff1d(dofs_W, bdofs_V))\n",
    "\n",
    "    # Assemble lhs and rhs matrices\n",
    "    A = assemble_matrix(lhs, restriction=(restriction, restriction))\n",
    "    A.assemble()\n",
    "    B = assemble_matrix(rhs, restriction=(restriction, restriction))\n",
    "    B.assemble()\n",
    "\n",
    "    # Solve\n",
    "    eps = SLEPc.EPS().create()\n",
    "    eps.setOperators(A, B)\n",
    "    eps.setProblemType(SLEPc.EPS.ProblemType.GNHEP)\n",
    "    eps.setDimensions(1, PETSc.DECIDE, PETSc.DECIDE)\n",
    "    eps.setWhichEigenpairs(SLEPc.EPS.Which.TARGET_REAL)\n",
    "    eps.setTarget(1.e-5)\n",
    "    eps.getST().setType(SLEPc.ST.Type.SINVERT)\n",
    "    eps.getST().getKSP().setType(\"preonly\")\n",
    "    eps.getST().getKSP().getPC().setType(\"lu\")\n",
    "    eps.getST().getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "    eps.solve()\n",
    "    assert eps.getConverged() >= 1\n",
    "\n",
    "    # Extract leading eigenvalue and eigenvector\n",
    "    vr = create_vector_block([(restriction.index_map, restriction.index_map_bs)])\n",
    "    vi = create_vector_block([(restriction.index_map, restriction.index_map_bs)])\n",
    "    eigv = eps.getEigenpair(0, vr, vi)\n",
    "    r, i = eigv.real, eigv.imag\n",
    "    assert abs(i) < 1.e-10\n",
    "    assert r > 0., \"r = \" + str(r) + \" is not positive\"\n",
    "    print(\"Inf-sup constant (monolithic): \", np.sqrt(r))\n",
    "\n",
    "    # Transform eigenvector into eigenfunction\n",
    "    r_fun = Function(W)\n",
    "    vr.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)\n",
    "    with r_fun.vector.localForm() as r_fun_local, VecSubVectorWrapper(vr, W.dofmap, restriction) as vr_wrapper:\n",
    "        r_fun_local[:] = vr_wrapper\n",
    "    (u_fun_1, u_fun_2) = (r_fun.sub(0).sub(0).collapse(), r_fun.sub(0).sub(1).collapse())\n",
    "    p_fun = r_fun.sub(1).collapse()\n",
    "    normalize(u_fun_1, u_fun_2, p_fun)\n",
    "\n",
    "    return (r, u_fun_1, u_fun_2, p_fun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(eig_m, u_fun_1_m, u_fun_2_m, p_fun_m) = run_monolithic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_scalar_field_plot(mesh, scalar_field, name):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    grid.point_arrays[name] = scalar_field.compute_point_values()\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u_fun_1_m, \"u1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u_fun_2_m, \"u2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p_fun_m, \"p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Block FEniCSx formulation using a two independent function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_block():\n",
    "    # Function spaces\n",
    "    V = FunctionSpace(mesh, V_element)\n",
    "    Q = FunctionSpace(mesh, Q_element)\n",
    "\n",
    "    # Test and trial functions\n",
    "    (v, q) = (TestFunction(V), TestFunction(Q))\n",
    "    (u, p) = (TrialFunction(V), TrialFunction(Q))\n",
    "\n",
    "    # Variational forms\n",
    "    lhs = [[inner(grad(u), grad(v)) * dx, - div(v) * p * dx],\n",
    "           [- div(u) * q * dx, None]]\n",
    "    rhs = [[None, None],\n",
    "           [None, - p * q * dx]]\n",
    "    rhs[0][0] = Constant(mesh, 0.) * inner(u, v) * dx\n",
    "\n",
    "    # Define restriction for DOFs associated to homogenous Dirichlet boundary conditions\n",
    "    dofs_V = np.arange(0, V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts)\n",
    "    bdofs_V = locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets)\n",
    "    dofs_Q = np.arange(0, Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts)\n",
    "    restriction_V = DofMapRestriction(V.dofmap, np.setdiff1d(dofs_V, bdofs_V))\n",
    "    restriction_Q = DofMapRestriction(Q.dofmap, dofs_Q)\n",
    "    restriction = [restriction_V, restriction_Q]\n",
    "\n",
    "    # Assemble lhs and rhs matrices\n",
    "    A = assemble_matrix_block(lhs, bcs=[], restriction=(restriction, restriction))\n",
    "    A.assemble()\n",
    "    B = assemble_matrix_block(rhs, bcs=[], restriction=(restriction, restriction))\n",
    "    B.assemble()\n",
    "\n",
    "    # Solve\n",
    "    eps = SLEPc.EPS().create()\n",
    "    eps.setOperators(A, B)\n",
    "    eps.setProblemType(SLEPc.EPS.ProblemType.GNHEP)\n",
    "    eps.setDimensions(1, PETSc.DECIDE, PETSc.DECIDE)\n",
    "    eps.setWhichEigenpairs(SLEPc.EPS.Which.TARGET_REAL)\n",
    "    eps.setTarget(1.e-5)\n",
    "    eps.getST().setType(SLEPc.ST.Type.SINVERT)\n",
    "    eps.getST().getKSP().setType(\"preonly\")\n",
    "    eps.getST().getKSP().getPC().setType(\"lu\")\n",
    "    eps.getST().getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "    eps.solve()\n",
    "    assert eps.getConverged() >= 1\n",
    "\n",
    "    # Extract leading eigenvalue and eigenvector\n",
    "    vr = create_vector_block([(restriction_.index_map, restriction_.index_map_bs) for restriction_ in restriction])\n",
    "    vi = create_vector_block([(restriction_.index_map, restriction_.index_map_bs) for restriction_ in restriction])\n",
    "    eigv = eps.getEigenpair(0, vr, vi)\n",
    "    r, i = eigv.real, eigv.imag\n",
    "    assert abs(i) < 1.e-10\n",
    "    assert r > 0., \"r = \" + str(r) + \" is not positive\"\n",
    "    print(\"Inf-sup constant (block): \", np.sqrt(r))\n",
    "\n",
    "    # Transform eigenvector into eigenfunction\n",
    "    (u_fun, p_fun) = (Function(V), Function(Q))\n",
    "    vr.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)\n",
    "    with BlockVecSubVectorWrapper(vr, [V.dofmap, Q.dofmap], restriction) as vr_wrapper:\n",
    "        for vr_wrapper_local, component in zip(vr_wrapper, (u_fun, p_fun)):\n",
    "            with component.vector.localForm() as component_local:\n",
    "                component_local[:] = vr_wrapper_local\n",
    "    (u_fun_1, u_fun_2) = (u_fun.sub(0).collapse(), u_fun.sub(1).collapse())\n",
    "    normalize(u_fun_1, u_fun_2, p_fun)\n",
    "\n",
    "    return (r, u_fun_1, u_fun_2, p_fun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(eig_b, u_fun_1_b, u_fun_2_b, p_fun_b) = run_block()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u_fun_1_b, \"u1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u_fun_2_b, \"u2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p_fun_b, \"p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error computation between standard and block formulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_error(eig_m, eig_b, u_fun_1_m, u_fun_1_b, u_fun_2_m, u_fun_2_b, p_fun_m, p_fun_b):\n",
    "    err_inf_sup = abs(np.sqrt(eig_b) - np.sqrt(eig_m)) / np.sqrt(eig_m)\n",
    "    print(\"Relative error for inf-sup constant equal to\", err_inf_sup)\n",
    "    assert np.isclose(err_inf_sup, 0., atol=1.e-8)\n",
    "    # Even after normalization, eigenfunctions may have different signs. Try both and assume that the correct\n",
    "    # error computation is the one for which the error is minimum\n",
    "    err_1_plus = u_fun_1_b + u_fun_1_m\n",
    "    err_2_plus = u_fun_2_b + u_fun_2_m\n",
    "    err_p_plus = p_fun_b + p_fun_m\n",
    "    err_1_minus = u_fun_1_b - u_fun_1_m\n",
    "    err_2_minus = u_fun_2_b - u_fun_2_m\n",
    "    err_p_minus = p_fun_b - p_fun_m\n",
    "    err_1_plus_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(err_1_plus), grad(err_1_plus)) * dx)))\n",
    "    err_2_plus_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(err_2_plus), grad(err_2_plus)) * dx)))\n",
    "    err_p_plus_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(err_p_plus * err_p_plus * dx)))\n",
    "    err_1_minus_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(err_1_minus), grad(err_1_minus)) * dx)))\n",
    "    err_2_minus_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(err_2_minus), grad(err_2_minus)) * dx)))\n",
    "    err_p_minus_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(err_p_minus * err_p_minus * dx)))\n",
    "    u_fun_1_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_fun_1_m), grad(u_fun_1_m)) * dx)))\n",
    "    u_fun_2_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_fun_2_m), grad(u_fun_2_m)) * dx)))\n",
    "    p_fun_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(p_fun_m * p_fun_m * dx)))\n",
    "\n",
    "    def select_error(err_plus, err_plus_norm, err_minus, err_minus_norm, vec_norm, component_name):\n",
    "        ratio_plus = err_plus_norm / vec_norm\n",
    "        ratio_minus = err_minus_norm / vec_norm\n",
    "        if ratio_minus < ratio_plus:\n",
    "            print(\"Relative error for \", component_name, \"component of eigenvector equal to\",\n",
    "                  ratio_minus, \"(the one with opposite sign was\", ratio_plus, \")\")\n",
    "            assert np.isclose(ratio_minus, 0., atol=1.e-6)\n",
    "        else:\n",
    "            print(\"Relative error for\", component_name, \"component of eigenvector equal to\",\n",
    "                  ratio_plus, \"(the one with opposite sign was\", ratio_minus, \")\")\n",
    "            assert np.isclose(ratio_plus, 0., atol=1.e-6)\n",
    "\n",
    "    select_error(err_1_plus, err_1_plus_norm, err_1_minus, err_1_minus_norm, u_fun_1_norm, \"velocity 1\")\n",
    "    select_error(err_2_plus, err_2_plus_norm, err_2_minus, err_2_minus_norm, u_fun_2_norm, \"velocity 2\")\n",
    "    select_error(err_p_plus, err_p_plus_norm, err_p_minus, err_p_minus_norm, p_fun_norm, \"pressure\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_error(eig_m, eig_b, u_fun_1_m, u_fun_1_b, u_fun_2_m, u_fun_2_b, p_fun_m, p_fun_b)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
